#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: scatter和gather用法.py
@time: 2020/9/23 16:03
@desc:
'''

from torch.nn.parallel import scatter,gather,replicate
import torch
import torch.nn as nn



def data_parallel(module, input, device_ids, output_device=None):
    if not device_ids:
        return module(input)

    if output_device is None:
        output_device = device_ids[0]

    replicas = nn.parallel.replicate(module, device_ids)
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    return nn.parallel.gather(outputs, output_device)

a = torch.randn(10,4,4)
device_ids = [0,1,2]
inputs = scatter(a,device_ids,dim=0)
print(inputs)
outputs = gather(inputs,0,dim=0)
print(outputs)